#! /usr/bin/python
# -*- coding: utf-8 -*-
# @Time    : 2017/11/3 11:32
# @Author  : Deyu.Tian
from osgeo import gdal, osr
import os
from util import *
import config
import numpy as np

def ascii2tif(fold, out):
    """
    ascii to geotiff convertion
    create the gdal output file as geotiff
    set the no data value
    set the geotransform
    numpy.genfromtxt('your file', numpy.int8) #looks like int from you example
    reshape your array to the shape you need
    write out the array.
    :param fname:
    :return:
    """
    txts = list_all_txts(fold)

    for txt in txts:
        # Set file vars
        output_file = "{}\\{}aaa.tif".format(out, txt[-11:-4])
        # Create gtif
        driver = gdal.GetDriverByName("GTiff")
        dst_ds = driver.Create(output_file, 321, 161, 1, gdal.GDT_CFloat32)
        raster = np.genfromtxt(txt, dtype=np.float32)
        raster = np.flipud(raster)
        print(raster.shape)

        # top left x, w-e pixel resolution, rotation, top left y, rotation, n-s pixel resolution
        dst_ds.SetGeoTransform([60, 0.25, 0, 55, 0, 0.25])

        # set the reference info
        srs = osr.SpatialReference()
        srs.SetWellKnownGeogCS("WGS84")
        dst_ds.SetProjection(srs.ExportToWkt())

        # write the band
        dst_ds.GetRasterBand(1).WriteArray(raster)
        break


def bands2Multispec(imgId, b):
    """
    bands to rgb picture using gdal_merge
    :param bands:
    :param out:
    :return:
    """
    if not os.path.exists("{}/{}".format(config.outDir, imgId)):
        cmd = "mkdir {}/{}".format(config.outDir, imgId)
        os.system(cmd)
        out = "{}/{}/{}_M.tif".format(config.outDir, imgId, imgId)
        cmd = "gdal_merge.py -separate -o {} -a_nodata -9999 " \
              "{} {} {} {} {} {} {} {} {} {}"\
            .format(out, b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7], b[8], b[9])
        os.system(cmd)


def bands2rgb(imgId, bands):
    """
    bands to rgb picture using gdal_merge
    :param bands:
    :param out:
    :return:
    """
    if not os.path.exists("{}/{}".format(config.outDir, imgId)):
        cmd = "mkdir {}/{}".format(config.outDir, imgId)
        os.system(cmd)
        out = "{}/{}/{}_rgb.tif".format(config.outDir, imgId, imgId)
        cmd = "gdal_merge.py -separate -o {} -a_nodata -9999 {} {} {}".format(out, bands[2], bands[1], bands[0])
        os.system(cmd)

def reprojection(geo_dem):
    """
    gdal reprojection
    :param dem:
    :return:
    """
    cmd = "gdalwarp -t_srs '+proj=utm +zone=43N +datum=WGS84' -dstnodata -9999 -overwrite {} {}_toUTM.tif".format(geo_dem, geo_dem[:-4])
    os.system(cmd)

def resampling(utm_dem):
    """
    down sampling dem data for better visualize
    :param utm_dem:
    :return:
    """
    cmd = "gdal_translate -tr 500 500 -r cubic -a_nodata 0 -stats {} {}_resamp.tif".format(utm_dem, utm_dem[:-4])
    os.system(cmd)

def resamplingAster(aster_v1):
    """
    down sampling dem data for better visualize
    :param utm_dem:
    :return:
    """
    cmd = "gdal_translate -tr 30 30 -r cubic -a_nodata -9999 -stats {} {}_resamp.tif".format(aster_v1, aster_v1[:-4])
    os.system(cmd)


def img2array(img):
    """

    read dems to array by gdal

    :param imgfn path of geotiff

    :return narray of geotiff

    """

    img_data = gdal.Open(img)
    img_array = img_data.ReadAsArray()

    return img_array


def dem2array(dem1, dem2):
    """

    read dems to array by gdal

    :param imgfn path of geotiff

    :return narray of geotiff

    """

    dem1_data = gdal.Open(dem1)
    dem1_array = dem1_data.ReadAsArray()

    dem2_data = gdal.Open(dem2)
    dem2_array = dem2_data.ReadAsArray()

    return dem1_array, dem2_array


def dem3array(dem1, dem2, dem3):
    """

    read dems to array by gdal

    :param imgfn path of geotiff

    :return narray of geotiff

    """

    dem1_data = gdal.Open(dem1)
    dem1_array = dem1_data.ReadAsArray()

    dem2_data = gdal.Open(dem2)
    dem2_array = dem2_data.ReadAsArray()

    dem3_data = gdal.Open(dem3)
    dem3_array = dem3_data.ReadAsArray()

    return dem1_array, dem2_array, dem3_array

def imgs2rgb():
    """
    converts all imgs in a folder to rgb
    :return:
    """
    subdirs = list_of_subdirectories(config.miniralDir)
    for img_folder in subdirs:
        imgId = img_folder.split('/')[-2]
        V1 = "{}/{}_V1.tif".format(img_folder, imgId)
        V2 = "{}/{}_V2.tif".format(img_folder, imgId)
        V3N = "{}/{}_V3N.tif".format(img_folder, imgId)
        V3B = "{}/{}_V3B.tif".format(img_folder, imgId)
        S4 = "{}/{}_S4.tif".format(img_folder, imgId)
        S5 = "{}/{}_S5.tif".format(img_folder, imgId)
        S6 = "{}/{}_S6.tif".format(img_folder, imgId)
        S7 = "{}/{}_S7.tif".format(img_folder, imgId)
        S8 = "{}/{}_S8.tif".format(img_folder, imgId)
        S9 = "{}/{}_S9.tif".format(img_folder, imgId)
        Bands = [V1, V2, V3N, V3B, S4, S5, S6, S7, S8, S9]
        bands2Multispec(imgId, Bands)


def stretch_N(bands, lower_percent=5, higher_percent=95):
    """

    将图像波段值缩放到0-1之间

    __author == n01z3:

    :param band: input image array

    : return: bands_copy: strcd bands

    """
    bands_copy = np.zeros_like(bands).astype(np.float32) #fix bug

    N = bands.shape[2]
    for i in range(N):
        oldmin = 0
        oldmax = 1
        newmin = np.percentile(bands[:, :, i], lower_percent)
        #print('5% percentile new min value of image band is{}'.format(newmin))

        newmax = np.percentile(bands[:, :, i], higher_percent)
        #print('95% percentile new max value of image band is{}'.format(newmax))

        strcd = oldmin + \
        (bands[:, :, i] - newmin)* (oldmax - oldmin) / (newmax - newmin)
        strcd[strcd < oldmin] = oldmin
        strcd[strcd > oldmax] = oldmax
        bands_copy[:, :, i] = strcd
    return bands_copy.astype(np.float32)

def readNDSI(imagefile):
    imgds = gdal.Open(imagefile)
    imggt = imgds.GetGeoTransform()
    print('raster geotransform coeffs:', imggt[0], imggt[1], imggt[2], imggt[3], imggt[4], imggt[5])
    band = imgds.GetRasterBand(1)
    b = band.ReadAsArray()
    return imggt, b

def readmeltsnow(imagefile):
    imgds = gdal.Open(imagefile)
    imggt = imgds.GetGeoTransform()
    print('raster geotransform coeffs:', imggt[0], imggt[1], imggt[2], imggt[3], imggt[4], imggt[5])
    band = imgds.GetRasterBand(1)
    b = band.ReadAsArray()
    return imggt, b

def read_tif_metadata(tifffile):
    """
    read tiff imggt
    :param tifffile:
    :return:
    """
    imgds = gdal.Open(tifffile)
    imggt = imgds.GetGeoTransform()
    print('raster geotransform coeffs:', imggt[0], imggt[1], imggt[2], imggt[3], imggt[4], imggt[5])
    band = imgds.GetRasterBand(1)
    b = band.ReadAsArray()
    return imggt, b

def array2rasterUTM(newRasterfn, panTransform, array):
    """

    :param newRasterfn:
    :param panTransform: imggt
    :param array:
    :return:
    """
    cols = array.shape[1]
    rows = array.shape[0]

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(newRasterfn, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform((panTransform[0], panTransform[1], panTransform[2], panTransform[3],
                               panTransform[4], panTransform[5]))
    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRasterSRS = osr.SpatialReference()
    outRasterSRS.ImportFromEPSG(32643)
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()

def array2rasterwgs84(newRasterfn, panTransform, array):
    """

    :param newRasterfn:
    :param panTransform: imggt
    :param array:
    :return:
    """
    cols = array.shape[1]
    rows = array.shape[0]

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(newRasterfn, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform((panTransform[0], panTransform[1], panTransform[2], panTransform[3],
                               panTransform[4], panTransform[5]))
    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRasterSRS = osr.SpatialReference()
    outRasterSRS.ImportFromEPSG(4326)
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()

if __name__ == '__main__':
    #reprojection(STRM_DEM)
    #bands2rgb(imgId, Bands)
    #resampling(STRM_DEM)
    #calcNDSI(imgId)
    #ThresholdNDSI(imgId)
    imgs2rgb()